from torch.utils.data import Dataset
from PIL import Image
import os
import torchvision.transforms.functional as TF
import matplotlib.pyplot as plt
from augument import horizontal_flip, vertical_flip
import numpy as np

#数据集处理类，加载图像和相应的掩码，并提供数据增强功能（如翻转）来增加数据多样性
class SegDatasets(Dataset):
    def __init__(self, image_size, data_root, input_dir_name,
                 label_dir_name, h_flip, v_flip, train=True):
        # 初始化函数，设置数据集的属性和参数
        self.image_size = image_size           # 图像尺寸
        self.data_dir = data_root              # 数据集的根目录，包含图像和相应标签
        self.input_dir_name = input_dir_name   # 原始图像的目录名称
        self.label_dir_name = label_dir_name   # 图像对应掩码、标签的目录名称
        self.h_flip = h_flip                   # 布尔值，表示是否要进行水平翻转的数据增强
        self.v_flip = v_flip                   # 布尔值，表示是否要进行垂直翻转的数据增强
        self.train = train                     # 布尔值，表示是否为训练模式
        if self.train:                         # 如果是训练模式，则设置前缀为train_
            self.prefix = "train_"
        else:                                  # 如果是验证模式，则设置前缀为val_
            self.prefix = "val_"

        # 检查数据集目录是否存在
        if not os.path.exists(self.data_dir):
            raise Exception(r"[!] data set does not exist!")

        # 获取所有训练数据的名称，存储到self.files列表中，并按照字母顺序排列
        self.files = sorted(os.listdir(os.path.join(self.data_dir,
                                                    self.prefix + self.input_dir_name)))

    def __getitem__(self, item):
        # 获取特定索引 item 处的样本名
        file_name = self.files[item]
        # 打开img和对应的mask，img是RGB图像，mask是灰度图像
        img = Image.open(os.path.join(self.data_dir,   #完整的图像文件路径
                                      self.prefix + self.input_dir_name,
                                      file_name)).convert('RGB')
        mask = Image.open(os.path.join(self.data_dir, #file_name[:-4] 是对掩码文件名进行切片操作，去掉末尾的后缀，添加 "_matte.png" 作为新后缀
                                       self.prefix + self.label_dir_name,
                                       file_name[:-4] + "_matte.png")).convert('L')

        # 将img和mask进行resize，统一为相同尺寸
        img = TF.resize(img, (self.image_size, self.image_size))
        mask = TF.resize(mask, (self.image_size, self.image_size))

        #训练模式执行数据增强
        if self.train:
            # 以0.5的概率进行数据增强，增强方式必须保证img和mask的变换是完全对应的
            if self.h_flip and np.random.random() > 0.5:
               img, mask = horizontal_flip(img, mask)

            if self.v_flip and np.random.random() > 0.5:
                img, mask = vertical_flip(img, mask)

        # 将图像和掩码转为tensor类型
        img = TF.to_tensor(img)
        mask = TF.to_tensor(mask)

        # 以字典形式返回img、mask和img的名字
        out = {'human': img, 'mask': mask, "img_name": file_name}
        return out

    def __len__(self):        
        # 返回数据集的长度，即样本数量
        return len(self.files)


if __name__ == "__main__":
    #参数和常量的设置
    BETA1 = 0.9                  #优化算法中的超参数
    BETA2 = 0.999                #优化算法中的超参数
    DATA_ROOT = "human_dataset"  #存储数据集的根目录
    INPUT_DIR_NAME = "human"     #数据集中存储输入图像的目录名称
    LABEL_DIR_NAME = "mask"      #数据集中存储标签（掩码）的目录名称
    LR = 0.0001                  #学习率，用于控制模型参数在训练过程中的更新步长
    BATCH_SIZE = 16              #训练时的批次大小
    H_FLIP = True                #布尔变量，是否进行水平翻转的数据增强
    V_FLIP = True                #布尔变量，是否进行垂直翻转的数据增强
    RESULTS_DIR = "results"      #存储训练结果的目录
    EPOCHS = 50                  #训练的轮数（遍历数据集的次数）
    IMAGE_SIZE = 224             #处理图像时的目标大小
    IMG_SAVE_FREQ = 100          #控制保存训练过程中图像的频率
    PTH_SAVE_FREQ = 2            #控制保存模型的频率

    VAL_BATCH_SIZE = 1           #验证过程的批次大小
    VAL_FREQ = 1                 #控制验证的频率（多少个周期后进行一次验证）
    # 构建训练Dataset，创建训练集对象
    train_set = SegDatasets(IMAGE_SIZE, DATA_ROOT, INPUT_DIR_NAME, LABEL_DIR_NAME, H_FLIP, V_FLIP, train=True)

    # 数据集中图像数量
    print("num of Train set {}".format(len(train_set)))

    # 获取数据集中的第4条数据的原始图像img、掩码mask和图像名称
    img = train_set[3]["human"]
    mask = train_set[3]["mask"]
    name = train_set[3]["img_name"]

    # 展示第4条数据的原始图像img和掩码mask
    plt.subplot(1, 2, 1)
    #将Tensor转换为numpy数组，并将张量的通道维度（CxHxW，其中 C 表示通道数，H 表示高度，W 表示宽度）转换为常用的图像维度（HxWxC）
    plt.imshow(np.transpose(img.numpy(), (1, 2, 0)))  # 调整img展示的方式
    plt.subplot(1, 2, 2)
    #将Tensor转换为numpy数组，并消除维度为 1 的通道维度，从 (1, H, W) 转换为 (H, W)
    plt.imshow(np.squeeze(mask.numpy(), axis=0), cmap='gray')  # 调整mask展示的方式
    plt.show()
